#!/usr/bin/env/ python
# -*- coding: utf-8 -*-
""" 
@author:tsl 
@file: checkout.py 
@time: 2020/04/13 
"""

import cv2,os

def colorShow(id):
    colors = {
        "0":(0,0,255),
        "1":(0,255,0),
        "2":(255,0,0),
        "3":(0,255,255),
        "4":(255,0,255),
        "5":(255,255,0),
        "6":(0,255,255),
        "7":(255,255,255),
        "8":(0,0,0)
    }

    if(str(id) not in list(list(colors.keys()) )):
        return (127,1127,127)

    return colors[str(id)]

def resultShow(src,dst,imgPath):
    for file in os.listdir(src):
        locPath = os.path.join(src,file)
        name = os.path.splitext(file)[0]
        img = cv2.imread(os.path.join(imgPath,name+".jpg"))
        img = cv2.resize(img,(608,608))

        with open(locPath,"r") as f:
            lines = f.readlines()
            for line in lines:
                if("useTime" not in line):
                    x, y, w, h, score, classID = line.rstrip().split(",")
                    x, y, w, h,score,  classID = map(int,[x, y, w, h,score, classID])
                    cv2.rectangle(img,(x,y),(x+w,y+h),colorShow(classID),2)
                    cv2.putText(img,str(classID)+","+str(score),(x,y),1,1,colorShow(classID),2)

                else:
                    useTime = line.rstrip()
                    cv2.putText(img,str(useTime),(10,20),2,1,(0,0,255),2)
        try:
            cv2.imwrite(os.path.join(dst,name+".jpg"),img)
            cv2.imshow("img",img)
            cv2.waitKey(1)

        except:
            continue

def main():
    # src = "boatSrc/result"
    # imgPath = "/mnt/data/3519/caffeCvtTools/nnie_bgr/data416"

    src = "../headSrc"
    imgPath = "/mnt/data/3519_/demo/mpp/sample/3559/sample_src/mpp/sample/svp/head/ins_nnie/data/nnie_image/head4"
    dst = "../headResult/head4"
    resultShow(src,dst,imgPath)

if __name__ == "__main__":
    main()